{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jJy_iUhoZJQw"
      },
      "source": [
        "# Welcome to Yuna Ai LLM trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Kncicf0oOVl"
      },
      "outputs": [],
      "source": [
        "!pip install --no-deps bitsandbytes accelerate xformers==0.0.29.post3 peft trl triton cut_cross_entropy unsloth_zoo\n",
        "!pip install sentencepiece protobuf \"datasets>=3.4.1\" huggingface_hub hf_transfer\n",
        "!pip install --no-deps unsloth\n",
        "!pip install --no-deps git+https://github.com/huggingface/transformers.git # Only for Gemma 3N\n",
        "!pip install --no-deps --upgrade timm # Only for Gemma 3N\n",
        "!pip install transformers==4.53.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PO8W1X7naZN4"
      },
      "source": [
        "# Initialize the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2eSvM9zX_2d3"
      },
      "outputs": [],
      "source": [
        "from unsloth import FastModel\n",
        "max_seq_length = 16384\n",
        "\n",
        "model, tokenizer = FastModel.from_pretrained(\n",
        "    model_name = \"modelpath\",\n",
        "    max_seq_length = max_seq_length,\n",
        "    full_finetuning = True, # use without LoRA\n",
        "    dtype = None,\n",
        "    load_in_4bit = True,\n",
        "    #load_in_8bit = True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model = FastLanguageModel.get_peft_model(\n",
        "    model,\n",
        "    r = 128, # 128 or 256\n",
        "    target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
        "                      \"gate_proj\", \"up_proj\", \"down_proj\",\n",
        "                      \"embed_tokens\", \"lm_head\"],\n",
        "    lora_alpha = 128, # 128 or 256\n",
        "    lora_dropout = 0,\n",
        "    bias = \"all\", # \"all\" or \"lora_only\"\n",
        "    use_gradient_checkpointing = \"unsloth\",\n",
        "    random_state = 8,\n",
        "    use_rslora = True,\n",
        "    loftq_config = None, # And LoftQ\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Freeze the vision tower and multi-modal projector\n",
        "for name, param in model.named_parameters():\n",
        "    if \"vision_tower\" in name or \"multi_modal_projector\" in name:\n",
        "        param.requires_grad_(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Bug Fixes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Add this line right here if doesn't work!\n",
        "model.resize_token_embeddings(262164)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sgz-KTv3adqp"
      },
      "source": [
        "# Training time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "95_Nn-89DhsL"
      },
      "outputs": [],
      "source": [
        "# Connect the stuff\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "dataset = load_dataset(\"json\", data_files=\"datapath.jsonl\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# This is where the magic happens\n",
        "from unsloth import UnslothTrainer, UnslothTrainingArguments\n",
        "\n",
        "trainer = UnslothTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    train_dataset=dataset[\"train\"],\n",
        "    eval_dataset = None,\n",
        "    dataset_text_field=\"text\",\n",
        "    max_seq_length=max_seq_length,\n",
        "    packing=False, # only for tiny datasets\n",
        "    args=UnslothTrainingArguments(\n",
        "        output_dir=\"OutputFolder\",\n",
        "        per_device_train_batch_size=1,\n",
        "        gradient_accumulation_steps=4,\n",
        "        num_train_epochs = 50,\n",
        "        warmup_ratio=0.1,\n",
        "        learning_rate=1e-5,\n",
        "        embedding_learning_rate=1e-6,  # 2-10x smaller than learning_rate\n",
        "        bf16=True,\n",
        "        logging_steps=1,\n",
        "        optim=\"adamw_8bit\",\n",
        "        weight_decay=0.01,\n",
        "        lr_scheduler_type=\"cosine\",\n",
        "        seed=8,\n",
        "        max_grad_norm=2.0,\n",
        "        report_to=\"tensorboard\",\n",
        "        save_strategy = \"epoch\",\n",
        "        save_steps=1,\n",
        "        gradient_checkpointing=True,\n",
        "        save_safetensors = False # safetensors are bad!\n",
        "    ),\n",
        ")\n",
        "\n",
        "trainer_stats = trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from transformers import Gemma3ForConditionalGeneration, AutoTokenizer\n",
        "import torch\n",
        "\n",
        "# Load the model and tokenizer\n",
        "model = Gemma3ForConditionalGeneration.from_pretrained(\"modelpath\", trust_remote_code=True, torch_dtype=torch.float16, device_map=\"auto\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"modelpath\", trust_remote_code=True)\n",
        "prompt = \"Once upon a time,\"  # Your input prompt\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
        "\n",
        "# Apply softmax temperature scaling to the logits before sampling\n",
        "with torch.no_grad():\n",
        "    outputs = model(**inputs, max_new_tokens=50, output_hidden_states=True, return_dict_in_generate=True)\n",
        "    logits = outputs.logits[:, -1, :]\n",
        "    temperature = 0.7\n",
        "    probs = torch.softmax(logits / temperature, dim=-1)\n",
        "    next_token = torch.multinomial(probs, num_samples=1)\n",
        "\n",
        "# Append the sampled token to the input IDs and continue generation\n",
        "input_ids = torch.cat([inputs[\"input_ids\"], next_token], dim=-1)\n",
        "generated_ids = model.generate(input_ids=input_ids, max_new_tokens=49)  # Adjust max_new_tokens accordingly\n",
        "\n",
        "generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=False)\n",
        "print(generated_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3adTid7VaRYy"
      },
      "source": [
        "# Export"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import Gemma3ForConditionalGeneration, AutoTokenizer\n",
        "from peft import PeftModel\n",
        "\n",
        "# Load the base model in FP16\n",
        "base_model = Gemma3ForConditionalGeneration.from_pretrained(\n",
        "    \"modelpath\",\n",
        "    low_cpu_mem_usage=True,\n",
        "    return_dict=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"cpu\",\n",
        "    tie_word_embeddings=False\n",
        ")\n",
        "\n",
        "base_model.lm_head.weight.data = base_model.get_input_embeddings().weight.data.clone()\n",
        "\n",
        "# Load the PEFT model\n",
        "model = PeftModel.from_pretrained(base_model, \"checkpointpath\")\n",
        "\n",
        "# Merge the LoRA weights into the base model\n",
        "merged_model = model.merge_and_unload()\n",
        "\n",
        "# Update the generation config\n",
        "merged_model.generation_config.temperature = 0.7\n",
        "merged_model.generation_config.top_p = 1.0\n",
        "merged_model.generation_config.do_sample = True\n",
        "\n",
        "# Load and configure the tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"checkpointpath\", trust_remote_code=True)\n",
        "tokenizer.padding_side = \"right\"\n",
        "\n",
        "# Save the full merged model and tokenizer\n",
        "merged_model.save_pretrained(\"outputpath\", max_shard_size=\"40GB\", safe_serialization=False)\n",
        "tokenizer.save_pretrained(\"outputpath\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wpvveZO4aSAa"
      },
      "outputs": [],
      "source": [
        "# check the folder size\n",
        "\n",
        "!du -sh folderpath"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIgmdRf5RgDz"
      },
      "outputs": [],
      "source": [
        "!huggingface-cli upload username/model \"modelfolder\" ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Create"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import torch\n",
        "import os\n",
        "\n",
        "# 1. Load the base model\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    \"modelpath\",\n",
        "    low_cpu_mem_usage=True,\n",
        "    torch_dtype=torch.bfloat16,\n",
        "    device_map=\"cuda\",\n",
        "    trust_remote_code=True\n",
        ")\n",
        "\n",
        "# 2. Initialize tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"modelpath\", trust_remote_code=True)\n",
        "\n",
        "# Retrieve existing special tokens\n",
        "existing_special_tokens = tokenizer.special_tokens_map.get(\"additional_special_tokens\", [])\n",
        "\n",
        "# Define new special tokens\n",
        "new_special_tokens = [\n",
        "    \"<memory>\", \"</memory>\", \"<shujinko>\", \"</shujinko>\",\n",
        "    \"<kanojo>\", \"</kanojo>\", \"<dialog>\", \"<yuki>\", \"</yuki>\",\n",
        "    \"<yuna>\", \"</yuna>\", \"<hito>\", \"</hito>\", \"<qt>\", \"</qt>\",\n",
        "    \"<action>\", \"</action>\", \"<data>\", \"</data>\"]\n",
        "\n",
        "# Combine existing and new special tokens, ensuring no duplicates\n",
        "all_special_tokens = list(set(existing_special_tokens + new_special_tokens))\n",
        "\n",
        "# Add them to the tokenizer\n",
        "tokenizer.add_special_tokens({\"additional_special_tokens\": all_special_tokens})\n",
        "\n",
        "print(f\"New vocab size: {len(tokenizer)}\")\n",
        "\n",
        "# 3. Resize model embeddings\n",
        "base_model.resize_token_embeddings(len(tokenizer))\n",
        "\n",
        "# 4. Save the updated model and tokenizer\n",
        "base_model.save_pretrained(\"new_model_dir\", max_shard_size=\"40GB\", safe_serialization=False)\n",
        "tokenizer.save_pretrained(\"new_model_dir\")\n",
        "\n",
        "print(\"New model with additional special tokens created and saved.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# HF conversion\n",
        "import torch\n",
        "from transformers import Gemma3ForConditionalGeneration, AutoTokenizer\n",
        "\n",
        "# Step 2: Load the PyTorch model (this automatically unties the weights)\n",
        "print(\"Reloading model to untie weights...\")\n",
        "untied_model = Gemma3ForConditionalGeneration.from_pretrained(\n",
        "    \"modelpath\",\n",
        "    low_cpu_mem_usage=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"cpu\",\n",
        ")\n",
        "# Step 1: Load the tokenizer\n",
        "print(\"Loading tokenizer...\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"modelpath\", use_fast=True)\n",
        "untied_model.lm_head.weight = torch.nn.Parameter(untied_model.language_model.embed_tokens.weight.clone())\n",
        "\n",
        "# Step 3: Save in safetensors format\n",
        "print(\"Converting to safetensors format...\")\n",
        "untied_model.save_pretrained(\"outputpath\", max_shard_size=\"40GB\", safe_serialization=True)\n",
        "tokenizer.save_pretrained(\"outputpath\")\n",
        "\n",
        "print(\"Conversion complete!\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
